#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <quadrotor_msgs/PositionCommand.h>
#include <std_srvs/SetBool.h>
#include <tf/tf.h>

mavros_msgs::State current_state;
geometry_msgs::PoseStamped local_pos;
quadrotor_msgs::PositionCommand latest_pos_cmd;
ros::Time last_cmd_time;

bool use_external_pos_cmd = false;
geometry_msgs::Point last_valid_pos;

const double TRAJ_TIMEOUT = 3.0; // 秒
const double JUMP_THRESHOLD = 2.0; // 米，xy跳变保护
const double HOVER_Z = 1.5;

// 回调函数
void state_cb(const mavros_msgs::State::ConstPtr& msg){
    current_state = *msg;
}

void local_pos_cb(const geometry_msgs::PoseStamped::ConstPtr& msg){
    local_pos = *msg;
}

void posCmdCallback(const quadrotor_msgs::PositionCommand::ConstPtr& msg) {
    if (msg->trajectory_flag == quadrotor_msgs::PositionCommand::TRAJECTORY_STATUS_READY) {
        double dx = msg->position.x - last_valid_pos.x;
        double dy = msg->position.y - last_valid_pos.y;
        if (sqrt(dx*dx + dy*dy) < JUMP_THRESHOLD) {
            latest_pos_cmd = *msg;
            last_valid_pos = msg->position;
            last_cmd_time = ros::Time::now();
            ROS_INFO_THROTTLE(1.0, "[OFFB_NODE] ACCEPTED pos_cmd: (%.2f, %.2f)", msg->position.x, msg->position.y);
        } else {
            ROS_WARN_THROTTLE(1.0, "[OFFB_NODE] REJECTED pos_cmd: jump too large (%.2f, %.2f)", dx, dy);
        }
    }
}

bool togglePosCmdCallback(std_srvs::SetBool::Request &req, std_srvs::SetBool::Response &res) {
    use_external_pos_cmd = req.data;
    res.success = true;
    res.message = use_external_pos_cmd ? "外部轨迹接管 ENABLED" : "外部轨迹接管 DISABLED";
    ROS_WARN("[OFFB_NODE] %s", res.message.c_str());
    return true;
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "offb_ego");
    ros::NodeHandle nh;

    ros::Subscriber state_sub = nh.subscribe("/mavros/state", 10, state_cb);
    ros::Subscriber local_pos_sub = nh.subscribe("/mavros/local_position/pose", 10, local_pos_cb);
    ros::Subscriber pos_cmd_sub = nh.subscribe("/drone_1_planning/pos_cmd", 10, posCmdCallback);
    ros::ServiceServer pos_cmd_toggle_srv = nh.advertiseService("/toggle_external_pos_cmd", togglePosCmdCallback);

    ros::Publisher local_pos_pub = nh.advertise<geometry_msgs::PoseStamped>("/mavros/setpoint_position/local", 10);
    ros::ServiceClient arming_client = nh.serviceClient<mavros_msgs::CommandBool>("/mavros/cmd/arming");
    ros::ServiceClient set_mode_client = nh.serviceClient<mavros_msgs::SetMode>("/mavros/set_mode");

    ros::Rate rate(20.0);

    while (ros::ok() && !current_state.connected) {
        ros::spinOnce();
        rate.sleep();
    }

    geometry_msgs::PoseStamped pose;
    pose.pose.position.x = 0;
    pose.pose.position.y = 0;
    pose.pose.position.z = HOVER_Z;

    for (int i = 100; ros::ok() && i > 0; --i) {
        local_pos_pub.publish(pose);
        ros::spinOnce();
        rate.sleep();
    }

    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";

    mavros_msgs::CommandBool arm_cmd;
    arm_cmd.request.value = true;

    ros::Time last_request = ros::Time::now();
    int step = 0;
    int hover_counter = 0;
    last_valid_pos.x = 0;
    last_valid_pos.y = 0;
    last_valid_pos.z = 0;
    last_cmd_time = ros::Time::now();

    while (ros::ok()) {
        if (current_state.mode != "OFFBOARD" &&
            (ros::Time::now() - last_request > ros::Duration(5.0))) {
            if (set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent) {
                ROS_INFO("Offboard enabled");
            }
            last_request = ros::Time::now();
        } else {
            if (!current_state.armed &&
                (ros::Time::now() - last_request > ros::Duration(5.0))) {
                if (arming_client.call(arm_cmd) && arm_cmd.response.success) {
                    ROS_INFO("Vehicle armed");
                }
                last_request = ros::Time::now();
            } else {
                switch (step) {
                case 0: // 起飞并悬停
                    pose.pose.position.x = 0;
                    pose.pose.position.y = 0;
                    pose.pose.position.z = HOVER_Z;
                    if (fabs(local_pos.pose.position.z - HOVER_Z) < 0.1) {
                        if (++hover_counter > 100) {
                            step = 1;
                            ROS_INFO("Step 1: 等待目标点...");
                        }
                    } else {
                        hover_counter = 0;
                    }
                    break;

                case 1: // 飞至目标点
                    if (use_external_pos_cmd &&
                        (ros::Time::now() - last_cmd_time < ros::Duration(TRAJ_TIMEOUT))) {
                        // 使用 pos_cmd 控制飞行
                        pose.pose.position.x = latest_pos_cmd.position.x;
                        pose.pose.position.y = latest_pos_cmd.position.y;
                        pose.pose.position.z = HOVER_Z;
                        if (fabs(local_pos.pose.position.x - pose.pose.position.x) < 0.3 &&
                            fabs(local_pos.pose.position.y - pose.pose.position.y) < 0.3) {
                            if (++hover_counter > 50) {
                                step = 2;
                                ROS_INFO("Step 2: 飞回起点");
                            }
                        } else {
                            hover_counter = 0;
                        }
                    } else {
                        ROS_WARN_THROTTLE(1.0, "等待有效轨迹指令...");
                    }
                    break;

                case 2: // 飞回起点
                    pose.pose.position.x = 0;
                    pose.pose.position.y = 0;
                    pose.pose.position.z = HOVER_Z;
                    if (fabs(local_pos.pose.position.x) < 0.2 && fabs(local_pos.pose.position.y) < 0.2) {
                        if (++hover_counter > 50) {
                            step = 3;
                            ROS_INFO("Step 3: 执行降落");
                        }
                    } else {
                        hover_counter = 0;
                    }
                    break;

                case 3: // 自动降落
                    offb_set_mode.request.custom_mode = "AUTO.LAND";
                    if (current_state.mode != "AUTO.LAND" &&
                        (ros::Time::now() - last_request > ros::Duration(5.0))) {
                        if (set_mode_client.call(offb_set_mode) && offb_set_mode.response.mode_sent) {
                            ROS_INFO("AUTO.LAND enabled");
                        }
                        last_request = ros::Time::now();
                    }
                    break;

                default:
                    break;
                }
            }
        }

        // 始终发布控制指令（轨迹或默认悬停）
        pose.header.stamp = ros::Time::now();
        pose.header.frame_id = "map";

        tf::Quaternion q = tf::createQuaternionFromYaw(0.0);
        quaternionTFToMsg(q, pose.pose.orientation);

        local_pos_pub.publish(pose);

        ros::spinOnce();
        rate.sleep();
    }

    return 0;
}
